# factors.py

from abc import ABC, abstractmethod
import pandas as pd
import numpy as np


class Factor(ABC):
    """
    因子抽象基类 (Abstract Base Class)。
    所有具体因子都必须继承此类并实现 calculate 方法。
    """

    @abstractmethod
    def calculate(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        根据输入的合并数据计算因子值。

        :param data: 包含K线和衍生数据（如volume_delta）的多重索引DataFrame。
        :return: 一个包含单列因子值的DataFrame，索引与输入数据对齐。
        """
        pass


class MovingAverageCrossFactor(Factor):
    """
    均线交叉因子。
    当快线从下向上穿过慢线时，产生做多信号 (1)。
    当快线从上向下穿过慢线时，产生做空信号 (-1)。
    """

    def __init__(self, short_window: int, long_window: int):
        if short_window >= long_window:
            raise ValueError("短周期窗口必须小于长周期窗口。")
        self.short_window = short_window
        self.long_window = long_window
        self.name = f"MACross_{short_window}_{long_window}"

    def calculate(self, data: pd.DataFrame) -> pd.DataFrame:
        """计算均线交叉信号。"""
        factor_data = pd.DataFrame(index=data.index)

        # 按币种分组计算移动平均线
        short_ma = data.groupby('symbol')['close'].transform(lambda x: x.rolling(window=self.short_window).mean())
        long_ma = data.groupby('symbol')['close'].transform(lambda x: x.rolling(window=self.long_window).mean())

        # 生成信号
        signal = pd.Series(np.nan, index=data.index)

        # 金叉: 快线上穿慢线
        golden_cross = (short_ma > long_ma) & (short_ma.shift(1) <= long_ma.shift(1))
        signal[golden_cross] = 1.0

        # 死叉: 快线下穿慢线
        death_cross = (short_ma < long_ma) & (short_ma.shift(1) >= long_ma.shift(1))
        signal[death_cross] = -1.0

        # 信号持续，直到出现反向信号
        signal = signal.groupby('symbol').transform(lambda x: x.ffill()).fillna(0)

        factor_data[self.name] = signal
        return factor_data


class VolumeDeltaFactor(Factor):
    """
    交易量增量因子。
    计算volume_delta的移动平均值，当其为正时产生做多倾向，为负时产生做空倾向。
    """

    def __init__(self, window: int):
        self.window = window
        self.name = f"VolumeDelta_{window}"

    def calculate(self, data: pd.DataFrame) -> pd.DataFrame:
        """计算平滑后的交易量增量信号。"""
        factor_data = pd.DataFrame(index=data.index)

        # 按币种分组计算 volume_delta 的移动平均
        smoothed_delta = data.groupby('symbol')['volume_delta'].transform(
            lambda x: x.rolling(window=self.window).mean())

        # 信号为1（做多倾向）或-1（做空倾向）
        signal = np.sign(smoothed_delta).fillna(0)

        factor_data[self.name] = signal
        return factor_data